import json
from fastapi import FastAPI
from models.mcp_handlers import get_mcp_servers, add_mcp_servers, \
    update_mcp_version, my_delete_mcp_server, my_get_mcp_server, my_update_mcp_server
import uvicorn
from server.schemas import MCPCreateSchema
from models.mcp import McpServer
from fastapi.exceptions import HTTPException
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
from mcp import StdioServerParameters
import settings
import os
import ast
import aiofiles
from fastapi import UploadFile
from models.file_handler import add_file,get_file_list,get_file,delete_files
from models.file import FileInfo
from utils.zlrag import ZLRagAction
from fastapi.middleware.cors import CORSMiddleware


app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"]
)

rag_action = ZLRagAction()

@app.get("/api/mcp")
async def get_mcp_server():
    mcp_servers = await get_mcp_servers()
    return {
        "mcps": [mcp.as_dict() for mcp in mcp_servers],
        "total": len(mcp_servers)
    }


@app.post("/api/mcp")
async def add_mcp_server(mcp_info: MCPCreateSchema):
    name = mcp_info.name
    transport = mcp_info.transport
    command = mcp_info.command
    args = mcp_info.args
    env = mcp_info.env
    headers = mcp_info.headers
    mcp_server = McpServer(name=name, transport=transport, command=command, args=args, env=env, headers=headers)
    await add_mcp_servers(mcp_server)
    await update_mcp_version()
    return {
        "mcp": mcp_server.as_dict()
    }


@app.delete("/api/mcp/{mcp_id}")
async def delete_mcp_server(mcp_id: str):
    await my_delete_mcp_server(mcp_id)
    await update_mcp_version()
    return "ok"


@app.post("/api/mcp/test/{mcp_id}")
async def mcp_test(mcp_id: int):
    mcp_server = await my_get_mcp_server(mcp_id)
    if not mcp_server:
        raise HTTPException(status_code=400, detail="mcp_server不存在")
    if mcp_server.transport == "streamable_http":
        try:
            async with streamablehttp_client(mcp_server.url) as (read, write, _):
                async with ClientSession(read, write) as session:
                    await session.initialize()
        except Exception as e:
            raise HTTPException(status_code=510, detail=str(e))

    elif mcp_server.transport == "sse":
        try:
            async with sse_client(mcp_server.url) as (read, write):
                async with ClientSession(read, write) as session:
                    await session.initialize()
        except Exception as e:
            raise HTTPException(status_code=510, detail=str(e))

    else:
        args = []
        env = {}
        try:
            if type(mcp_server.args) == str:
                args = ast.literal_eval(mcp_server.args)
            elif type(mcp_server.args) == list:
                args = mcp_server.args
        except Exception as e:
            raise HTTPException(status_code=510, detail="args解析失败")
        try:
            if mcp_server.env is not None:
                if type(mcp_server.env) == str:
                    env = ast.literal_eval(mcp_server.env)
                else:
                    env = mcp_server.env
        except Exception as e:
            raise HTTPException(status_code=510, detail="env解析失败")

        server_params = StdioServerParameters(
            command=mcp_server.command,
            args=args,
            env=env
        )
        try:
            async with stdio_client(server_params) as (read, write):
                async with ClientSession(read, write) as session:
                    await session.initialize()
        except Exception as e:
            raise HTTPException(status_code=510, detail=str(e))

    await my_update_mcp_server(mcp_id, True)
    return "ok"

@app.post("/api/file/upload")
async def upload_file(file: UploadFile):
    file_path = os.path.join(settings.FILE_UPLOAD_DIR, file.filename)
    async with aiofiles.open(file_path, "wb") as fp:
        await fp.write(await file.read())

    await rag_action.file_to_vector(file_path)
    file_info = FileInfo(file_path=file_path)
    await add_file(file_info)
    return "ok"

@app.get("/api/file")
async def file_list():
    result = await get_file_list()
    return {
        "files": [file.as_dict() for file in result],
        "total": len(result)
    }

@app.delete("/api/file/{file_id}")
async def delete_file(file_id: int):
    file_info = await get_file(file_id)
    if not file_info:
        raise HTTPException(status_code=400, detail="文件不存在")
    try:
        os.remove(file_info.file_path)
    except Exception as e:
        pass

    await rag_action.delete_file_vector(file_info.file_path)
    await delete_files(file_id)
    return "ok"

if __name__ == '__main__':
    uvicorn.run(app, host="0.0.0.0", port=8000)
